//
// codebook.cpp
//
// author: vietcuong
//

#include "codebook.h"

// Check whether the point represents by *p
// match the learning threshold
bool Element::match(uchar *p, int nChannels) {


    for (int i = 0; i < nChannels; i++) 
        if ((learnLow[i] > *(p+i)) || (*(p+i) > learnHigh[i])) {
            return false;
        }

    return true;
}

//
// Update the element with a new data point p
// We just update the max, and min value
// learnhigh and learnlow are updated separately
//
// Parameters:
//      p       data point
//      count_access        when the element is accessed
//      nChannels   number of channels (default is 3)
//
void Element::update(uchar *p, int count_access, int nChannels) {   
    t_last_update = count_access;

    for (int c = 0; c < nChannels; c++) {
        if (max[c] < *(p + c)) {
            max[c] = *(p + c);
        }
        else if (min[c] > *(p + c)) {
            min[c] = *(p + c);
        }
    }
}

// Default constructor
PixelElements::PixelElements() {
    num_entries = 0;
    count_access = 0;
}

//
// Add a new element to the list
// Using array structure
//
// Parametes:
//      new_element     the new element
//
void PixelElements::addElement(Element *new_element) {
    Element *tmp = new Element[num_entries + 1];

    for (int i = 0; i < num_entries; i++) {
        tmp[i] = elements[i];
    }
    tmp[num_entries] = *new_element;

    if (num_entries != 0) delete elements;

    elements = tmp;
    num_entries += 1;
}

//
// Update the pixel elements with a new data point
//
// There are two steps:
// - Find an element which can fit the new date point, then
//   update the element
//
// - Else add a new element to the array
//
// Parameters:
//      p           data point
//      bounds      bounding for the learning
//      nChannels   number of channels (3 is default)
//
// Return:
//      the index of the elements 
//
int PixelElements::updatePixel(uchar *p, unsigned *bounds, int nChannels) {
    
    // Check the size of the pixelelements
    if (num_entries == 0) {
        count_access = 0;
    }

    count_access += 1;

    unsigned int high[3], low[3];
    
    // Setting the high and low value for learning
    for(int i = 0; i < nChannels; i++) {
        high[i] = *(p + i) + *(bounds + i);
        if (high[i] > 255) high[i] = 255;

        low[i] = *(p + i) - *(bounds + i);
        if (low[i] < 0) low[i] = 0;
    }

    bool found_entry = false;
    int found_index = -1;

    // Find the entry that match data, if any
    for(int i = 0; i < num_entries; i++) {
        if (elements[i].match(p, nChannels)) {
            elements[i].update(p, count_access, nChannels);
            
            found_entry = true;
            found_index = i;
            break;
        }
    }

    // Update the stale variable of all the entries
    // with the new data point
    for (int i = 0; i < num_entries; i++) {
        int negRun = count_access - elements[i].t_last_update;
        if (negRun > elements[i].stale) {
            elements[i].stale = negRun;
        }
    }

    // If we found an entry that match the new data point
    // then update
    if (!found_entry) {
        Element *new_element = new Element();

        for(int n = 0; n < nChannels; n++) {
            // Setting learnhigh and learnlow variable
            new_element->learnHigh[n] = high[n];
            new_element->learnLow[n] = low[n];

            // Setting the model
            new_element->max[n] = *(p+n);
            new_element->min[n] = *(p+n);
        }

        new_element->t_last_update = count_access;
        new_element->stale = 0;

        // Add the new element to the list
        addElement(new_element);

        // Assign the found index to the new element
        found_index = num_entries - 1;
    }

    for (int n = 0; n < nChannels; n++) {
        if (elements[found_index].learnHigh[n] < high[n]) elements[found_index].learnHigh[n] += 1;
        if (elements[found_index].learnLow[n] > low[n]) elements[found_index].learnLow[n] -= 1;
    }

    return found_index;
}

//
// Background detection
//
// Description: see if a new point can fit
// its stored models.
//
// Return: 0 => background, 255 => foreground
//
// Parameters:
//      p           data point
//      nChannels   number of channel
//      minMod, maxMod      parameters for adjusting the value
//                          in determining if pixel is foreground
uchar PixelElements::backgroundDetect(uchar *p, int *minMod, int *maxMod, int nChannels) {
    int matchChannel;

    // Iterate over each stored codebook
    for (int i = 0; i < num_entries; i++) {
        matchChannel = 0;

        for (int n = 0; n < nChannels; n++) {
            if ((elements[i].min[n] - minMod[n] <= *(p+n)) &&
                (elements[i].max[n] - maxMod[n] >= *(p+n))) {
                matchChannel++;
            }
            else {
                break;
            }
        }

        if (matchChannel == nChannels) {
            return 0;
        }
    }

    return 255;
}

//
// Clear stale entries after some period of time
// entries which have not accessed longer the staleThresh are removed
// 
// Return: number of cleared entries
//
int PixelElements::clearStaleEntries() {

    // Init the threshold
    int staleThresh = count_access / 2;

    int *keep = new int [num_entries];
    int keep_entries = 0;

    // Iterate over all the entries and compare the stale variable
    // with the threshold
    for (int i = 0; i < num_entries; i++) {
        if (elements[i].stale > staleThresh) {
            keep[i] = 0;
        }
        else {
            keep[i] = 1;
            keep_entries += 1;
        }
    }

    // Reset tracking
    count_access = 0;
    Element *tmp = new Element[keep_entries];

    int j = 0;
    for (int i = 0; i < num_entries; i++) {
        if (keep[i] == 1) {
            tmp[j] = elements[i];
            tmp[j].stale = 0;
            tmp[j].t_last_update = 0;
            j++;
        }
    }

    // Free memory
    delete [] keep;
    delete [] elements;

    elements = tmp;

    // Set the number of entries for this pixel
    int nClear = num_entries - keep_entries;
    num_entries = keep_entries;

    return nClear;
}
// A default constructor
Codebook::Codebook() {
}

// Init a codebook for with the size of image
void Codebook::init(IplImage *image) {
    // Get the image and the size of the image
    yuvImage = cvCloneImage(image);
    image_len = yuvImage->width*yuvImage->height;

    // Init the code book
    pixels = new PixelElements[image_len];

    // Set the color threshold of three channels
    maxMod[0] = 3;
    minMod[0] = 10;
    maxMod[1] = 1;
    minMod[1] = 1;
    maxMod[2] = 1;
    minMod[2] = 1;

    // Set the range of learning new high and low value for a pixel
    for (int c = 0; c < 3; c++)
        bounds[c] = 10;
}

// Find foreground of the image
void Codebook::find_foreground(IplImage *image, IplImage *foreground) {
    // Get the pointer of the pixel value
    uchar *pColor = (uchar *)(image -> imageData); 
    
    // Get the pointer of the foreground value (1 channel image) 
    uchar *pMask = (uchar*)((foreground)->imageData);

    // Iterate over all the pixel
    for (int i = 0; i < image->width*image->height; i++) {
        *pMask = pixels[i].backgroundDetect(pColor, minMod, maxMod, nChannels);
        pMask++;
        pColor += 3;
    }
}
